{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 2, Topic 3: Attacking Across MixColumns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *In Part 2, Topic 2, you saw how to apply CPA attacks to a hardware implementaiton of AES. As it turns out, this attack will not work against many hardware implementations of AES. This is thanks to the fact that the implementation doesn't need to store the output ciphertext in the same register as the state. In this case, the nearest distance for the state register will usually be between the output of the 9th round and the output of the 8th round. Thanks to MixColumns, our leakage model will no longer depend on a single byte at a time, making our CPA attack completely infeasible.*\n",
    "\n",
    "*In this lab, we'll examine and run a modified CPA attack, allowing our attack to succeed even with a MixColumns operation.*\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "* Understanding how the modified CPA attack allows attacks across MixColumns\n",
    "* Using ChipWhisperer/SCARED to analyze traces\n",
    "\n",
    "**NOTE:** This lab requires [SCARED](https://github.com/eshard/scared), which may be difficult to install on some systems (noteably Windows). If you're unable to install SCARED, you'll need to skip this lab.\n",
    "\n",
    "Targets known to work with this leakage model:\n",
    "\n",
    "* STM32F4\n",
    "* STM32L4\n",
    "* STM32L5\n",
    "* K82F MMCAU\n",
    "* CW305 Default AES Implementation\n",
    "* MBEDTLS and TINYAES128C\n",
    "\n",
    "A target with hardware AES is recommended as they often require this attack instead of a simpler attack; however, software AES is also vulnerable to this attack (though no more than a normal CPA attack)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start with, let's take a look at an example AES implementation that the last round state diff we used in Lab 2_2 won't work for:\n",
    "\n",
    "![](img/AES_MixCol.png)\n",
    "\n",
    "AES operations are in orange, while registers are in purple. The mux before AddRoundKey allows the plaintext to be loaded at the beginning, as well as MixColumns to be skipped during the last round. As we can see, the demux after AddRoundKey stops the ciphertext from being put into the state register. Instead, if we work backwards from the ciphertext, we can see the two closest states will be the output of the 9th round and the output of the 8th round. The big issue that we run into is that MixColumns, as the name suggests, combines the 4 bytes an AES state column together. Up to that point, each byte was completely independent of one another, meaning we could also evaluate each byte independently, turning a $2^{128}$ search into 16 $2^8$ searches. Now, with MixColumns, we need to evaluate 4 bytes at a time, giving us 4 $2^{32}$ searches. \n",
    "\n",
    "Attacking from the plaintext side is no better - the MixColumns is still there. Does that mean if we cross a MixColumns in our leakage model, our attack is thwarted? Luckily, the answer to this is no! Instead of doing a full random plaintext like our usual attack, let's instead make all the bytes except for a column or row constant and we'll walk through the start of AES:\n",
    "\n",
    "$\n",
    "\\left[\\begin{array}{llll}\n",
    "v_0 & v_1 & v_2 & v_3 \\\\\n",
    "c_0 & c_3 & c_6 & c_9 \\\\\n",
    "c_1 & c_4 & c_7 & c_{10} \\\\\n",
    "c_2 & c_5 & c_8 & c_{11}\n",
    "\\end{array}\\right]$\n",
    "\n",
    "Thanks to shift rows, that will give us 3 constant bytes and 1 variable byte in each column (everything is $'$ here because of SubBytes):\n",
    "\n",
    "$\n",
    "\\left[\\begin{array}{llll}\n",
    "v'_0 & v'_1 & v'_2 & v'_3 \\\\\n",
    "c'_3 & c'_6 & c'_9 & c'_0 \\\\\n",
    "c'_7 & c'_{10} & c'_1 & c'_{4} \\\\\n",
    "c'_{11} & c'_2 & c'_5 & c'_{8}\n",
    "\\end{array}\\right]$\n",
    "\n",
    "MixColumns (on a single column) is the following operation ($v$ is the variable byte, $c_n$ are the constants). Note that the input of MixColumns is really just the output of SubBytes, our usual attack point, since ShiftRows doesn't modify the values, it just moves them:\n",
    "\n",
    "$\\left[\\begin{array}{l}\n",
    "d_{0} \\\\\n",
    "d_{1} \\\\\n",
    "d_{2} \\\\\n",
    "d_{3}\n",
    "\\end{array}\\right]=\\left[\\begin{array}{llll}\n",
    "2 & 3 & 1 & 1 \\\\\n",
    "1 & 2 & 3 & 1 \\\\\n",
    "1 & 1 & 2 & 3 \\\\\n",
    "3 & 1 & 1 & 2\n",
    "\\end{array}\\right]\\left[\\begin{array}{l}\n",
    "v'_0 \\\\\n",
    "c'_{3} \\\\\n",
    "c'_{7} \\\\\n",
    "c'_{11}\n",
    "\\end{array}\\right]$\n",
    "\n",
    "$d_0$, for example, will be:\n",
    "\n",
    "$d_0 = 2v'_0 + 3c'_3 + c'_7 + c'_{11}$\n",
    "\n",
    "We can combine the three constants into a single constant:\n",
    "\n",
    "$d_0 = 2v'_0 + c_a$\n",
    "\n",
    "It might not be immediately obvious, but this has actually taken our $2^{32}$ attack down to $2^{16}$! If $c_a$ wasn't there, this would pretty much just be our regular CPA attack, so doing a CPA attack for each possible value of $c_a$ should allow us to to recover the key! This constant we've combined everything into is actually extremely powerful: any constants XOR'd ($+$) with with the SBox output will completely disappear, meaning the AddRoundKey following MixColumns is also eliminated! It is due to this fact that this attack will work against many different target implementations, including software AES, STM32F4 style implementations (which was vulnerable to last round state diff), and many models resistant to last round state diff.\n",
    "\n",
    "While this seems very promising, it does suffer from a few issues:\n",
    "\n",
    "1. $2^{16}$ is still a fairly large search space, meaning our attack will take a very long time\n",
    "1. $c$ is only being XOR'd here, which means we'll run into ghost peaks\n",
    "1. This XOR also means the correct and incorrect values for $c$ will have very similar correlations.\n",
    "\n",
    "Instead, we can attack $v'$ one bit at a time. This is very similar to a single bit DPA attack, except we're still using correlation instead of the difference. Now the constant will only have the effect of inverting the correlation (if that constant bit is 1) or not inverting it (if that constant bit is 0). This means we can completely remove the effect of the constant by taking the absolute value of the correlation, taking us back to a $2^8$ attack! We also don't have to worry about the ghost peaks or the similar correlations for $c_a$.\n",
    "\n",
    "As you might expect, on its own this attack works very poorly - 1 bit is not enough information for a CPA attack. Much better is to perform the attack on each bit of $2v$ and sum the absolute correlations. This isn't the only spot $v$ shows up though! We also have the other $d$ values:\n",
    "\n",
    "$d_1 = v'_0 + 2c'_3 + 3c'_7 + c'_{11} = v'_0 + c_b$\n",
    "\n",
    "$d_2 = v'_0 + c'_3 + 2c'_7 + 3c'_{11} = v'_0 + c_c$\n",
    "\n",
    "$d_3 = 3v'_0 + c'_3 + c'_7 + 2c'_{11} = 3v'_0 + c_d$\n",
    "\n",
    "Adding in these brings us up to 32 CPA attacks. If we only needed Hamming weight, this would be pretty simple, with our leakage model being:\n",
    "\n",
    "`h = ((2*sbox(pt0 + key0)>>bit) & 0x01`\n",
    "\n",
    "Things get a bit trickier with the Hamming distance. The key can be incorporated into $c$ since it's constant. However, we need to line up the MixColumns output with the correct byte in the input. This is still simple with the top byte of each column ($d_0$ lines up with $p_0$, $d_1$ lines up with $p_1$, and so on), but later bytes get shifted by ShiftRows. For example, attacking $k_1$, $d_0$ will line up with $p_{12}$, $d_1$ will line up with $p_{13}$ and so on. They also have a different order in MixColumns - $p_0$ goes 2, 1, 1, 3, but $p_1$ goes 3, 2, 1, 1.\n",
    "\n",
    "Aside from being more complicated than our usual CPA attack, this attack also has the disadvantage of only targeting 4 key bytes per capture campaign, since the other bytes in the AES column have to be constant. It's also a fairly lengthy analysis since we need to do 256 CPA attacks to recover all 16 key bytes.\n",
    "\n",
    "As an aside, if we didn't need to go through ARK1, we could make a big optimization on the attack. We can actually calculate the constant values by using the sign of each correlation. These constants can then be used to calculate the corresponding key from a system of equations. This version of the attack is detailed in https://eprint.iacr.org/2019/343.pdf. This paper also details an optimization that would allow us to bring the attack down to a single capture campaign, as well as extensions to AES-192 and AES-256, though we won't take a look at either of these extensions in this lab.\n",
    "\n",
    "This attack was originally detailed in https://eprint.iacr.org/2016/249.pdf by Amir Moradi and Tobias Schneider. Many thanks to them for the theory behind this tutorial!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "scope = cw.scope()\n",
    "target = cw.target(scope)\n",
    "scope.default_setup()\n",
    "\n",
    "## target specific setup after here..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'MBEDTLS'\n",
    "VERSION='HARDWARE'\n",
    "SS_VER='SS_VER_2_0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../../Setup_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\" \"$SS_VER\"\n",
    "cd ../../../firmware/mcu/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2 SS_VER=$3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = '../../../firmware/mcu/simpleserial-aes/simpleserial-aes-{}.hex'.format(PLATFORM)\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capture\n",
    "\n",
    "For capture, we need to do four separate capture campaigns to recover the key, since we're only targeting 4 bytes at a time. ChipWhisperer includes a special key-text pair to change out which parts of the plaintext are constant and which are variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import trange\n",
    "import cwtvla\n",
    "import numpy as np\n",
    "vec_type = 'row'\n",
    "ktp = cw.ktp.VarVec(vec_type=vec_type)\n",
    "key, pt = ktp.next()\n",
    "scope.adc.samples = 10000\n",
    "N = 500\n",
    "projects = []\n",
    "\n",
    "target.simpleserial_write('k', key)\n",
    "target.simpleserial_wait_ack()\n",
    "\n",
    "for cmpgn in trange(4):\n",
    "    project = cw.create_project(f\"Var_Vec_{cmpgn}\", overwrite=True)\n",
    "    projects.append(project)\n",
    "    for i in trange(N, leave=False):\n",
    "        ktp.var_vec = cmpgn\n",
    "        key, text = ktp.next()\n",
    "        trace = cw.capture_trace(scope, target, text, key)\n",
    "        if trace is None:\n",
    "            continue\n",
    "        project.traces.append(trace)\n",
    "    project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "for i in range(5):\n",
    "    plt.plot(projects[0].waves[i])\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis\n",
    "\n",
    "If you closed the project, reopen it now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "projects = []\n",
    "for i in range(4):\n",
    "    project = cw.open_project(f\"Var_Vec_{i}\")\n",
    "    projects.append(project)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ChipWhisperer/SCARED will take care all of the analysis for us here. You can adjust the number of traces you use with `n_traces`. Windowing can also be very important with this attack. You can adjust the window you analyzer with `trace_slice`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chipwhisperer.analyzer.attacks.attack_mix_columns import AttackMixColumns\n",
    "attack = AttackMixColumns(projects, vec_type=vec_type, hd=False)\n",
    "results = attack.run(n_traces=100, trace_slice=slice(200, 1200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bytearray(results[\"guess\"]))\n",
    "print(bytearray(projects[0].keys[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the correlations of the best guess/key or second best guess/key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "plt.figure()\n",
    "for i in range(16):\n",
    "    c = results[\"corr\"][i]\n",
    "    maxes = np.max(c, axis=1)\n",
    "    guess = np.argsort(maxes)[-1]\n",
    "    guess2 = np.argsort(maxes)[-2]\n",
    "    actual = projects[0].keys[0][i]\n",
    "    x = np.argmax(c[actual])\n",
    "    if guess != actual:\n",
    "        plt.plot(c[guess], \"g-\")\n",
    "    else:\n",
    "        plt.plot(c[guess2], \"g-\")\n",
    "    plt.plot(c[actual], \"r--\")\n",
    "    plt.plot(x, c[actual][x], \"ro\")\n",
    "    print(f\"Best guess {hex(guess)} (corr={maxes[guess]}), next best = {maxes[guess2]}, real = {maxes[actual]}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Original Implementations\n",
    "\n",
    "Provided for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import scared\n",
    "import numpy as np\n",
    "b = None\n",
    "gal2=np.array((\n",
    "0x00,0x02,0x04,0x06,0x08,0x0a,0x0c,0x0e,0x10,0x12,0x14,0x16,0x18,0x1a,0x1c,0x1e,\n",
    "0x20,0x22,0x24,0x26,0x28,0x2a,0x2c,0x2e,0x30,0x32,0x34,0x36,0x38,0x3a,0x3c,0x3e,\n",
    "0x40,0x42,0x44,0x46,0x48,0x4a,0x4c,0x4e,0x50,0x52,0x54,0x56,0x58,0x5a,0x5c,0x5e,\n",
    "0x60,0x62,0x64,0x66,0x68,0x6a,0x6c,0x6e,0x70,0x72,0x74,0x76,0x78,0x7a,0x7c,0x7e,\n",
    "0x80,0x82,0x84,0x86,0x88,0x8a,0x8c,0x8e,0x90,0x92,0x94,0x96,0x98,0x9a,0x9c,0x9e,\n",
    "0xa0,0xa2,0xa4,0xa6,0xa8,0xaa,0xac,0xae,0xb0,0xb2,0xb4,0xb6,0xb8,0xba,0xbc,0xbe,\n",
    "0xc0,0xc2,0xc4,0xc6,0xc8,0xca,0xcc,0xce,0xd0,0xd2,0xd4,0xd6,0xd8,0xda,0xdc,0xde,\n",
    "0xe0,0xe2,0xe4,0xe6,0xe8,0xea,0xec,0xee,0xf0,0xf2,0xf4,0xf6,0xf8,0xfa,0xfc,0xfe,\n",
    "0x1b,0x19,0x1f,0x1d,0x13,0x11,0x17,0x15,0x0b,0x09,0x0f,0x0d,0x03,0x01,0x07,0x05,\n",
    "0x3b,0x39,0x3f,0x3d,0x33,0x31,0x37,0x35,0x2b,0x29,0x2f,0x2d,0x23,0x21,0x27,0x25,\n",
    "0x5b,0x59,0x5f,0x5d,0x53,0x51,0x57,0x55,0x4b,0x49,0x4f,0x4d,0x43,0x41,0x47,0x45,\n",
    "0x7b,0x79,0x7f,0x7d,0x73,0x71,0x77,0x75,0x6b,0x69,0x6f,0x6d,0x63,0x61,0x67,0x65,\n",
    "0x9b,0x99,0x9f,0x9d,0x93,0x91,0x97,0x95,0x8b,0x89,0x8f,0x8d,0x83,0x81,0x87,0x85,\n",
    "0xbb,0xb9,0xbf,0xbd,0xb3,0xb1,0xb7,0xb5,0xab,0xa9,0xaf,0xad,0xa3,0xa1,0xa7,0xa5,\n",
    "0xdb,0xd9,0xdf,0xdd,0xd3,0xd1,0xd7,0xd5,0xcb,0xc9,0xcf,0xcd,0xc3,0xc1,0xc7,0xc5,\n",
    "0xfb,0xf9,0xff,0xfd,0xf3,0xf1,0xf7,0xf5,0xeb,0xe9,0xef,0xed,0xe3,0xe1,0xe7,0xe5), dtype='uint8')\n",
    "\n",
    "gal3=np.array((\n",
    "0x00,0x03,0x06,0x05,0x0c,0x0f,0x0a,0x09,0x18,0x1b,0x1e,0x1d,0x14,0x17,0x12,0x11,\n",
    "0x30,0x33,0x36,0x35,0x3c,0x3f,0x3a,0x39,0x28,0x2b,0x2e,0x2d,0x24,0x27,0x22,0x21,\n",
    "0x60,0x63,0x66,0x65,0x6c,0x6f,0x6a,0x69,0x78,0x7b,0x7e,0x7d,0x74,0x77,0x72,0x71,\n",
    "0x50,0x53,0x56,0x55,0x5c,0x5f,0x5a,0x59,0x48,0x4b,0x4e,0x4d,0x44,0x47,0x42,0x41,\n",
    "0xc0,0xc3,0xc6,0xc5,0xcc,0xcf,0xca,0xc9,0xd8,0xdb,0xde,0xdd,0xd4,0xd7,0xd2,0xd1,\n",
    "0xf0,0xf3,0xf6,0xf5,0xfc,0xff,0xfa,0xf9,0xe8,0xeb,0xee,0xed,0xe4,0xe7,0xe2,0xe1,\n",
    "0xa0,0xa3,0xa6,0xa5,0xac,0xaf,0xaa,0xa9,0xb8,0xbb,0xbe,0xbd,0xb4,0xb7,0xb2,0xb1,\n",
    "0x90,0x93,0x96,0x95,0x9c,0x9f,0x9a,0x99,0x88,0x8b,0x8e,0x8d,0x84,0x87,0x82,0x81,\n",
    "0x9b,0x98,0x9d,0x9e,0x97,0x94,0x91,0x92,0x83,0x80,0x85,0x86,0x8f,0x8c,0x89,0x8a,\n",
    "0xab,0xa8,0xad,0xae,0xa7,0xa4,0xa1,0xa2,0xb3,0xb0,0xb5,0xb6,0xbf,0xbc,0xb9,0xba,\n",
    "0xfb,0xf8,0xfd,0xfe,0xf7,0xf4,0xf1,0xf2,0xe3,0xe0,0xe5,0xe6,0xef,0xec,0xe9,0xea,\n",
    "0xcb,0xc8,0xcd,0xce,0xc7,0xc4,0xc1,0xc2,0xd3,0xd0,0xd5,0xd6,0xdf,0xdc,0xd9,0xda,\n",
    "0x5b,0x58,0x5d,0x5e,0x57,0x54,0x51,0x52,0x43,0x40,0x45,0x46,0x4f,0x4c,0x49,0x4a,\n",
    "0x6b,0x68,0x6d,0x6e,0x67,0x64,0x61,0x62,0x73,0x70,0x75,0x76,0x7f,0x7c,0x79,0x7a,\n",
    "0x3b,0x38,0x3d,0x3e,0x37,0x34,0x31,0x32,0x23,0x20,0x25,0x26,0x2f,0x2c,0x29,0x2a,\n",
    "0x0b,0x08,0x0d,0x0e,0x07,0x04,0x01,0x02,0x13,0x10,0x15,0x16,0x1f,0x1c,0x19,0x1a), dtype='uint8')\n",
    "\n",
    "\n",
    "\n",
    "w=0x00\n",
    "lut_input = [[0, 1, 2, 3],\n",
    "        [4, 5, 6, 7],\n",
    "        [8, 9, 10, 11],\n",
    "        [12, 13, 14, 15]]\n",
    "\n",
    "def round_gen_0(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[0, 13, 10, 7],\n",
    "                     [4, 1, 14, 11],\n",
    "                     [8, 5, 2, 15],\n",
    "                     [12, 9, 6, 3]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(0, gal2[scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess))])\n",
    "    return res\n",
    "\n",
    "def round_gen_1(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[1, 14, 11, 4],\n",
    "                     [5, 2, 15, 8],\n",
    "                     [9, 6, 3, 12],\n",
    "                     [13, 10, 7, 0]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(0, scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess)))\n",
    "    return res\n",
    "\n",
    "def round_gen_2(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[2, 15, 8, 5],\n",
    "                     [6, 3, 12, 9],\n",
    "                     [10, 7, 0, 13],\n",
    "                     [14, 1, 4, 1]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(0, scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess)))\n",
    "    return res\n",
    "\n",
    "def round_gen_3(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[3, 12, 9, 6],\n",
    "                     [7, 0, 13, 10],\n",
    "                     [11, 4, 14, 1],\n",
    "                     [15, 8, 5, 2]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(0, gal3[scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess))])\n",
    "    return res\n",
    "\n",
    "leakage_cmpgns = []\n",
    "for campaign in range(4):\n",
    "    leakage_cmpgns.append([scared.attack_selection_function(lambda plaintext, guesses: round_gen_0(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_1(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_2(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_3(plaintext, guesses, campaign))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm.autonotebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import estraces, scared\n",
    "import numpy as np\n",
    "from tqdm.notebook import trange\n",
    "#del cw_traces\n",
    "campaign = 0\n",
    "n_traces = 10000\n",
    "key_guess = []\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "    \n",
    "\n",
    "\n",
    "for campaign in trange(0,4):\n",
    "    b = None\n",
    "    cw_traces = estraces.read_ths_from_ram(np.array(projects[campaign].waves)[:n_traces,:], \n",
    "                                      plaintext=np.array([textin for textin in projects[campaign].textins], dtype='uint8')[:n_traces])\n",
    "    #cw_traces = estraces.read_ths_from_ram(np.array(projects[campaign].waves)[:n_traces,550:900], \n",
    "    #                                  plaintext=np.array([textin for textin in projects[campaign].textins], dtype='uint8')[:n_traces])\n",
    "    for t in trange(4, leave=False):\n",
    "        for i in trange(8, leave=False):\n",
    "            container = scared.Container(cw_traces)\n",
    "            a = scared.CPAAttack(selection_function=leakage_cmpgns[campaign][t],\n",
    "                                model=scared.Monobit(i),\n",
    "                                discriminant=scared.maxabs)\n",
    "\n",
    "\n",
    "            a.run(container)\n",
    "            if b is None:\n",
    "                b = abs(a.results)\n",
    "            else:\n",
    "                b += abs(a.results)\n",
    "    for i in range(0+4*campaign, 4+4*campaign):\n",
    "        c = np.nan_to_num(b[:,i,:])\n",
    "        maxes = np.max(c, axis=1)\n",
    "        guess = np.argsort(maxes)[-1]\n",
    "        guess2 = np.argsort(maxes)[-2]\n",
    "        actual = projects[0].keys[0][i]\n",
    "        x = np.argmax(c[actual])\n",
    "        if guess != actual:\n",
    "            plt.plot(c[guess], \"g-\")\n",
    "        else:\n",
    "            plt.plot(c[guess2], \"g-\")\n",
    "        plt.plot(c[actual], \"r--\")\n",
    "        plt.plot(x, c[actual][x], \"ro\")\n",
    "        print(f\"Best guess {hex(guess)} (corr={maxes[guess]}), next best = {maxes[guess2]}, real = {maxes[actual]}\")\n",
    "        key_guess.append(guess)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bytearray(key_guess))\n",
    "print(bytearray(projects[0].keys[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.plot(c[0x16])\n",
    "plt.plot(c[0x7f])\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "for i in range(256):\n",
    "    plt.plot(c[i])\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import scared\n",
    "import numpy as np\n",
    "gal2=np.array((\n",
    "0x00,0x02,0x04,0x06,0x08,0x0a,0x0c,0x0e,0x10,0x12,0x14,0x16,0x18,0x1a,0x1c,0x1e,\n",
    "0x20,0x22,0x24,0x26,0x28,0x2a,0x2c,0x2e,0x30,0x32,0x34,0x36,0x38,0x3a,0x3c,0x3e,\n",
    "0x40,0x42,0x44,0x46,0x48,0x4a,0x4c,0x4e,0x50,0x52,0x54,0x56,0x58,0x5a,0x5c,0x5e,\n",
    "0x60,0x62,0x64,0x66,0x68,0x6a,0x6c,0x6e,0x70,0x72,0x74,0x76,0x78,0x7a,0x7c,0x7e,\n",
    "0x80,0x82,0x84,0x86,0x88,0x8a,0x8c,0x8e,0x90,0x92,0x94,0x96,0x98,0x9a,0x9c,0x9e,\n",
    "0xa0,0xa2,0xa4,0xa6,0xa8,0xaa,0xac,0xae,0xb0,0xb2,0xb4,0xb6,0xb8,0xba,0xbc,0xbe,\n",
    "0xc0,0xc2,0xc4,0xc6,0xc8,0xca,0xcc,0xce,0xd0,0xd2,0xd4,0xd6,0xd8,0xda,0xdc,0xde,\n",
    "0xe0,0xe2,0xe4,0xe6,0xe8,0xea,0xec,0xee,0xf0,0xf2,0xf4,0xf6,0xf8,0xfa,0xfc,0xfe,\n",
    "0x1b,0x19,0x1f,0x1d,0x13,0x11,0x17,0x15,0x0b,0x09,0x0f,0x0d,0x03,0x01,0x07,0x05,\n",
    "0x3b,0x39,0x3f,0x3d,0x33,0x31,0x37,0x35,0x2b,0x29,0x2f,0x2d,0x23,0x21,0x27,0x25,\n",
    "0x5b,0x59,0x5f,0x5d,0x53,0x51,0x57,0x55,0x4b,0x49,0x4f,0x4d,0x43,0x41,0x47,0x45,\n",
    "0x7b,0x79,0x7f,0x7d,0x73,0x71,0x77,0x75,0x6b,0x69,0x6f,0x6d,0x63,0x61,0x67,0x65,\n",
    "0x9b,0x99,0x9f,0x9d,0x93,0x91,0x97,0x95,0x8b,0x89,0x8f,0x8d,0x83,0x81,0x87,0x85,\n",
    "0xbb,0xb9,0xbf,0xbd,0xb3,0xb1,0xb7,0xb5,0xab,0xa9,0xaf,0xad,0xa3,0xa1,0xa7,0xa5,\n",
    "0xdb,0xd9,0xdf,0xdd,0xd3,0xd1,0xd7,0xd5,0xcb,0xc9,0xcf,0xcd,0xc3,0xc1,0xc7,0xc5,\n",
    "0xfb,0xf9,0xff,0xfd,0xf3,0xf1,0xf7,0xf5,0xeb,0xe9,0xef,0xed,0xe3,0xe1,0xe7,0xe5), dtype='uint8')\n",
    "\n",
    "gal3=np.array((\n",
    "0x00,0x03,0x06,0x05,0x0c,0x0f,0x0a,0x09,0x18,0x1b,0x1e,0x1d,0x14,0x17,0x12,0x11,\n",
    "0x30,0x33,0x36,0x35,0x3c,0x3f,0x3a,0x39,0x28,0x2b,0x2e,0x2d,0x24,0x27,0x22,0x21,\n",
    "0x60,0x63,0x66,0x65,0x6c,0x6f,0x6a,0x69,0x78,0x7b,0x7e,0x7d,0x74,0x77,0x72,0x71,\n",
    "0x50,0x53,0x56,0x55,0x5c,0x5f,0x5a,0x59,0x48,0x4b,0x4e,0x4d,0x44,0x47,0x42,0x41,\n",
    "0xc0,0xc3,0xc6,0xc5,0xcc,0xcf,0xca,0xc9,0xd8,0xdb,0xde,0xdd,0xd4,0xd7,0xd2,0xd1,\n",
    "0xf0,0xf3,0xf6,0xf5,0xfc,0xff,0xfa,0xf9,0xe8,0xeb,0xee,0xed,0xe4,0xe7,0xe2,0xe1,\n",
    "0xa0,0xa3,0xa6,0xa5,0xac,0xaf,0xaa,0xa9,0xb8,0xbb,0xbe,0xbd,0xb4,0xb7,0xb2,0xb1,\n",
    "0x90,0x93,0x96,0x95,0x9c,0x9f,0x9a,0x99,0x88,0x8b,0x8e,0x8d,0x84,0x87,0x82,0x81,\n",
    "0x9b,0x98,0x9d,0x9e,0x97,0x94,0x91,0x92,0x83,0x80,0x85,0x86,0x8f,0x8c,0x89,0x8a,\n",
    "0xab,0xa8,0xad,0xae,0xa7,0xa4,0xa1,0xa2,0xb3,0xb0,0xb5,0xb6,0xbf,0xbc,0xb9,0xba,\n",
    "0xfb,0xf8,0xfd,0xfe,0xf7,0xf4,0xf1,0xf2,0xe3,0xe0,0xe5,0xe6,0xef,0xec,0xe9,0xea,\n",
    "0xcb,0xc8,0xcd,0xce,0xc7,0xc4,0xc1,0xc2,0xd3,0xd0,0xd5,0xd6,0xdf,0xdc,0xd9,0xda,\n",
    "0x5b,0x58,0x5d,0x5e,0x57,0x54,0x51,0x52,0x43,0x40,0x45,0x46,0x4f,0x4c,0x49,0x4a,\n",
    "0x6b,0x68,0x6d,0x6e,0x67,0x64,0x61,0x62,0x73,0x70,0x75,0x76,0x7f,0x7c,0x79,0x7a,\n",
    "0x3b,0x38,0x3d,0x3e,0x37,0x34,0x31,0x32,0x23,0x20,0x25,0x26,0x2f,0x2c,0x29,0x2a,\n",
    "0x0b,0x08,0x0d,0x0e,0x07,0x04,0x01,0x02,0x13,0x10,0x15,0x16,0x1f,0x1c,0x19,0x1a), dtype='uint8')\n",
    "\n",
    "\n",
    "\n",
    "w=0x00\n",
    "lut_input = [[0, 1, 2, 3],\n",
    "        [4, 5, 6, 7],\n",
    "        [8, 9, 10, 11],\n",
    "        [12, 13, 14, 15]]\n",
    "\n",
    "def round_gen_0(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[0, 4, 8, 12],\n",
    "                     [13, 1, 5, 9],\n",
    "                     [10, 14, 2, 6],\n",
    "                     [7, 11, 15, 3]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(new_pt, gal2[scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess))])\n",
    "    return res\n",
    "\n",
    "def round_gen_1(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[1, 5, 9, 13],\n",
    "                     [14, 2, 6, 10],\n",
    "                     [11, 15, 3, 7],\n",
    "                     [4, 8, 12, 0]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(new_pt, scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess)))\n",
    "    return res\n",
    "\n",
    "def round_gen_2(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[2, 6, 10, 14],\n",
    "                     [15, 3, 7, 11],\n",
    "                     [8, 12, 0, 4],\n",
    "                     [5, 9, 13, 1]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(new_pt, scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess)))\n",
    "    return res\n",
    "\n",
    "def round_gen_3(plaintext, guesses, cmpgn):\n",
    "    lut_mix_column = [[3, 7, 11, 15],\n",
    "                     [12, 0, 4, 8],\n",
    "                     [9, 13, 1, 5],\n",
    "                     [6, 10, 14, 2]] # lut to find which pt to xor with mixcolumn output\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    new_pt = np.repeat(plaintext[:,lut_mix_column[cmpgn][0]][:, np.newaxis], 16, axis=1)\n",
    "    new_pt[:,lut_input[cmpgn][1]] = plaintext[:,lut_mix_column[cmpgn][1]]\n",
    "    new_pt[:,lut_input[cmpgn][2]] = plaintext[:,lut_mix_column[cmpgn][2]]\n",
    "    new_pt[:,lut_input[cmpgn][3]] = plaintext[:,lut_mix_column[cmpgn][3]]\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(new_pt, gal3[scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess))])\n",
    "    return res\n",
    "\n",
    "leakage_cmpgns = []\n",
    "for campaign in range(4):\n",
    "    leakage_cmpgns.append([scared.attack_selection_function(lambda plaintext, guesses: round_gen_0(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_1(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_2(plaintext, guesses, campaign)),\n",
    "    scared.attack_selection_function(lambda plaintext, guesses: round_gen_3(plaintext, guesses, campaign))])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
